import torch
import numpy as np
import cv2 as cv
import os
from torch.utils.data import DataLoader, Dataset

class garbageDataset(Dataset):
    def __init__(self, folderPath, size):
        super().__init__()
        self.size = size
        self.PicNames, self.PicLabels = [], []

        for folder in [folderPath + name for name in os.listdir(folderPath)]:
            self.PicNames.extend([folder + "/" + name for name in os.listdir(folder)])
            print(folder)
            self.PicLabels.extend([int(folder.split("/")[2]) for i in range(len(os.listdir(folder)))])

    def __getitem__(self, item):
        img = cv.imread(self.PicNames[item])
        img = cv.resize(img, (self.size, self.size))
        label = self.PicLabels[item]
        return torch.from_numpy(img).reshape(3, 224, 224).float(), torch.tensor(label).long()

    def __len__(self):
       return len(self.PicNames)

#test
if __name__ == "__main__":
    data = DataLoader(garbageDataset("./Data/trainData/"), num_workers=4, shuffle=True)

    for idx, (x, y) in enumerate(data):
        print(idx)
        print(x.shape)
        print(y.shape)
